#include <sstream>

template<class T, class U>
struct is_same_underlying : std::false_type {};

template<class T>
struct is_same_underlying<T,T> : std::true_type {};

template<class T>
struct is_same_underlying<std::add_lvalue_reference_t<T>, T> : std::true_type{};

template<class T>
struct is_same_underlying<std::add_const_t<T>, T> : std::true_type{};

template<class T>
struct is_same_underlying<std::add_lvalue_reference_t<std::add_const_t<T>>, T> :
	std::true_type {};

template< class T, class U >
inline constexpr bool is_same_underlying_v = is_same_underlying<T, U>::value;

/* call as getMPI_Datatype<yourType>(); */
template<typename T>
constexpr MPI_Datatype getMPI_Datatype(){
	if constexpr ( is_same_underlying_v<T, char> )
		return MPI_CHAR;
	if constexpr ( is_same_underlying_v<T, int> )
		return MPI_INT;
	if constexpr ( is_same_underlying_v<T, long int> )
		return MPI_LONG;
	if constexpr ( is_same_underlying_v<T, long long int> )
		return MPI_LONG_LONG;
	if constexpr ( is_same_underlying_v<T, unsigned int> )
		return MPI_UNSIGNED;
	if constexpr ( is_same_underlying_v<T, unsigned long> )
		return MPI_UNSIGNED_LONG;
	if constexpr ( is_same_underlying_v<T, unsigned long long> )
		return MPI_UNSIGNED_LONG_LONG;
	if constexpr ( is_same_underlying_v<T, float> )
		return MPI_FLOAT;
	if constexpr ( is_same_underlying_v<T, double> )
		return MPI_DOUBLE;
	/* add more types as needed */
}



namespace mpi_H {
	template<typename T>
	std::ostream& extract( std::ostream& out, const T& t) {
		return out << t;
	}

	template<typename T, typename ...Ts>
	std::ostream& extract( std::ostream& out, const T& t, Ts&& ... rest ) {
		extract(out, t);
		return extract(out, std::forward<Ts>(rest) ...);
	}
}

template<typename ... Ts>
void orderedPrint( const MPI_Comm_H& comm, Ts&& ... messages ){
	for (int i{0}; i<comm.size(); ++i){
		MPI_Barrier(comm);
		if (comm.rank() == i){
			std::stringstream ss;
			ss << "Rank ";
			ss.width(2);
			ss << comm.rank() << ": ";
			mpi_H::extract(ss, std::forward<Ts>(messages) ...) << "\n";
			std::cout << ss.str();
		}
	}
}

namespace mpiDetail {

template<typename Arr, typename T>
constexpr inline void underlyingTypeAssert(){
	if constexpr ( std::is_base_of_v<Eigen::DenseBase<Arr>, Arr> )
		static_assert( std::is_same_v<typename Arr::Scalar, T>,
"If the data points you supply do not match the array's underlying \
type, provide both template parameters to force a cast!"
		);
}

} // end namespace mpiDetail


/**
 * @brief Gathers single values into an array on root.
 * 
 * @tparam T 
 * @tparam Arr
 * @param datapoint 
 * @param comm 
 * @return ArrayTypeIfFundamental<T,Arr> 
 */
template<typename Arr, typename T>
ArrayTypeIfFundamental<Arr,T> MPI_Gather_H(
	T datapoint,
	const MPI_Comm_H& comm,
	int root
){
	mpiDetail::underlyingTypeAssert<Arr,T>();
	
	Arr arr;
	if ( comm.rank() == root ) arr.resize( comm.size() );

	MPI_Gather(
		&datapoint, 1, getMPI_Datatype<T>(),
		arr.data(), 1, getMPI_Datatype<T>(),
		root, comm
	);
	return arr;
}

/**
 * @brief Gathers equal length vector-like arrays onto concatenated vector-like
 * array of same type on root.
 * Array type must provide
 * - data pointer function called data(),
 * - length function called size()
 * 
 * @tparam Arr The array type
 * @param localData 
 * @param comm 
 * @param root 
 * @return ArrayIfNotFundamental
 */
template<typename Arr>
ArrayTypeIfNotFundamental<Arr> MPI_Gather_H(
	const Arr& localData,
	const MPI_Comm_H& comm,
	int root
){
	Arr globalArray;
	if ( comm.rank() == root )
		globalArray.resize( comm.size() * localData.size() );

	constexpr MPI_Datatype mpiType { 
		getMPI_Datatype<decltype(*( globalArray.data() ) )>()
	};

	MPI_Gather(
		localData.data(), localData.size(), mpiType,
		globalArray.data(), localData.size(), mpiType,
		root, comm
	);
	return globalArray;
}

template<typename Arr, typename T>
ArrayTypeIfFundamental<Arr,T> MPI_Allgather_H(
	T datapoint,
	const MPI_Comm_H& comm
){
	mpiDetail::underlyingTypeAssert<Arr,T>();
	
	Arr arr;
	arr.resize( comm.size() );

	MPI_Allgather(
		&datapoint, 1, getMPI_Datatype<T>(),
		arr.data(), 1, getMPI_Datatype<T>(),
		comm
	);
	return arr;
}

template<typename Arr>
ArrayTypeIfNotFundamental<Arr> MPI_Allgather_H(
	Arr localData,
	const MPI_Comm_H& comm
){
	Arr arr;
	arr.resize( comm.size() * localData.size() );

	constexpr MPI_Datatype type {
		getMPI_Datatype<decltype(*( localData.data() ) )>()
	};

	MPI_Allgather(
		localData.data(), localData.size(), type,
		arr.data(), localData.size(), type,
		comm
	);
	return arr;
}



/**
 * @brief Performs the necessary steps to invoke MPI_Gather_v, especially 
 * creating the recvcount and displacement arrays.
 * 
 * @tparam T Type of both the local array as well as the concatenated one
 * @param localArray 
 * @param comm 
 * @return T 
 */
template<typename T>
T MPI_Gatherv_H(
	const T& localArray,
	// MPI_Datatype mpiType,
	const MPI_Comm_H& comm,
	int root
){
	static_assert( !std::is_fundamental_v<T>,
		"For fundamental types, use MPI_Gather_H."
	);

	Eigen::ArrayXi sizes { MPI_Gather_H<Eigen::ArrayXi,int>(
		static_cast<int>( localArray.size() ), comm, root
	) };

	T globalArray;
	Eigen::ArrayXi displacements;
	if ( comm.rank() == root ){
		globalArray.resize( sizes.sum() );
		mpiDetail::calcDisplacements( displacements, sizes, comm.size() );
	}

	constexpr MPI_Datatype mpiType { 
		getMPI_Datatype<decltype(*( globalArray.data() ) )>()
	};
	
	MPI_Gatherv(
		localArray.data(), localArray.size(), mpiType,
		globalArray.data(), sizes.data(), displacements.data(), mpiType,
		root, comm
	);

	return globalArray;
}


template<typename T>
T MPI_Allgatherv_H(
	const T& localArray,
	const MPI_Comm_H& comm
){
	static_assert( !std::is_fundamental_v<T>,
		"For fundamental types, use MPI_Allgather_H."
	);

	Eigen::ArrayXi sizes { MPI_Allgather_H<Eigen::ArrayXi,int>(
		static_cast<int>( localArray.size() ), comm
	) };

	T globalArray;
	Eigen::ArrayXi displacements;
	globalArray.resize( sizes.sum() );
	mpiDetail::calcDisplacements( displacements, sizes, comm.size() );

	constexpr MPI_Datatype mpiType { 
		getMPI_Datatype<decltype(*( globalArray.data() ) )>()
	};
	
	MPI_Allgatherv(
		localArray.data(), localArray.size(), mpiType,
		globalArray.data(), sizes.data(), displacements.data(), mpiType,
		comm
	);

	return globalArray;
}



template<typename T>
std::enable_if_t<std::is_fundamental_v<T>,void> MPI_Bcast_H(
	T& t,
	const MPI_Comm_H& comm,
	int root
){
	MPI_Bcast( &t, 1, getMPI_Datatype<T>(), root, comm );
}

template<typename T>
std::enable_if_t<!std::is_fundamental_v<T>,void> MPI_Bcast_H(
	T& t,
	const MPI_Comm_H& comm,
	int root
){
	decltype(t.size()) size;
	if ( comm.rank() == root ) size = t.size();
	MPI_Bcast_H(size, comm, root);

	if ( comm.rank() != root )
		t.resize(size);
	
	MPI_Bcast(
		t.data(), size,
		getMPI_Datatype<decltype( *(t.data()) )>(),
		root, comm
	);
}

template<typename Derived>
void MPI_Bcast_H(
	Eigen::DenseBase<Derived>& d,
	const MPI_Comm_H& comm,
	int root,
	bool resize
){
	if (resize){
		int rows, cols;
		if ( comm.rank() == root ){
			rows = d.rows();
			cols = d.cols();
		}
		MPI_Bcast_H(rows, comm, root);
		MPI_Bcast_H(cols, comm, root);
		

		if ( comm.rank() != root )
			d.derived().resize(rows, cols);
	}

	MPI_Bcast(
		d.derived().data(), d.size(),
		getMPI_Datatype<typename Derived::Scalar>(),
		root, comm
	);
}

/* disallow implicit type conversion between int and bool */
template<typename T, typename U>
void MPI_Bcast_H( T&, const MPI_Comm_H&, U ) = delete;
template<typename Derived, typename T, typename U>
void MPI_Bcast_H( Eigen::DenseBase<Derived>&, const MPI_Comm_H&, T, U) = delete;


template<typename T>
std::enable_if_t<std::is_fundamental_v<T>, void> MPI_Send_H(
	const T& t, const MPI_Comm_H& comm, int to, int tag
){
	MPI_Send( &t, 1, getMPI_Datatype<T>(), to, tag, comm );
}

template<typename T>
std::enable_if_t<!std::is_fundamental_v<T>, void> MPI_Send_H(
	const T& t, const MPI_Comm_H& comm, int to, int tag
){
	MPI_Send_H( t.size(), comm, to, tag );
	MPI_Send( t.data(), t.size(),
		getMPI_Datatype<std::remove_pointer_t<decltype( t.data() )>>(),
		to, tag, comm );
}

template<typename T>
std::enable_if_t<std::is_fundamental_v<T>, T> MPI_Recv_H(
	const MPI_Comm_H& comm, int from,
	int tag, MPI_Status* status
){
	T t;
	MPI_Recv( &t, 1, getMPI_Datatype<T>(),
		from, tag, comm, status);
	return t;
}

template<typename T>
std::enable_if_t<!std::is_fundamental_v<T>, T> MPI_Recv_H(
	const MPI_Comm_H& comm, int from,
	int tag, MPI_Status* status
){
	T t;
	t.resize( MPI_Recv_H<decltype(t.size())>(comm, from, tag, status) );
	MPI_Recv( t.data(), t.size(),
		getMPI_Datatype<std::remove_pointer_t<decltype( t.data() )>>(),
		from, tag, comm, status);
	return t;
}